import sys
from PyQt5.QtWidgets import \
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,\
    QPushButton, QLineEdit, QFrame, QGridLayout,QComboBox,QLabel, QSlider,QMessageBox
from PyQt5.QtWidgets import QPlainTextEdit
from PyQt5.QtGui import QFontMetrics
from PyQt5.QtCore import Qt
from PyQt5 import QtWidgets, QtGui
from api import *
from PyQt5.QtGui import QPixmap, QImage
from PIL import Image
import os
from wr import *
from page2 import *


class Page1(QMainWindow):
    #get_pic = 0
    #image_path = 'example.png'
    def __init__(self,parent=None):
        super(Page1, self).__init__(parent)
        self.setWindowTitle('stable diffusion水印系统')
        self.setGeometry(100, 100, 1000, 800)
    ##
        # 创建中心部件  
        central_widget = QWidget(self)  
        self.setCentralWidget(central_widget)  
  
        # 创建2x2网格布局  
        grid_layout = QGridLayout(central_widget)  
        grid_layout.setContentsMargins(10, 10, 10, 10)  # 设置边距为10
        grid_layout.setSpacing(10)  # 设置控件之间的间距  
  
        # 创建四个模块并添加到网格布局中  
        module1 = self.create_module_1()
        module2 = self.create_module_2()
        module3 = self.create_module_3()
        module4 = self.create_module_4()
  
        # 设置行和列的拉伸比例  
        grid_layout.setRowStretch(0, 1)  # 第一行总宽度拉伸比例为6
        grid_layout.setRowStretch(0, 1)  # 第二行总宽度拉伸比例为6
  
        # 添加模块到布局中  
        grid_layout.addWidget(module1, 0, 0, 1, 1)  # 第一行占两列，宽度比例3:2
        grid_layout.addWidget(module2, 0, 1, 1, 1)  # 第一行占两列，宽度比例3:2
        grid_layout.addWidget(module3, 1, 0, 1, 1)  # 第二行占一列，宽度比例1:1
        grid_layout.addWidget(module4, 1, 1, 1, 1)  # 第二行占一列，宽度比例1:1
  
        self.setStyleSheet("""  
                    QWidget {  
                background-color: #f0f0f0;  
            }  
            
            QLabel {  
                border: 2px solid #555;  
                border-radius: 5px;  
                padding: 10px;  
                font-size: 12pt;  
                font-family: "宋体", sans-serif; /* 设置统一的字体 */  
                color: #000; /* 设置统一的文本颜色 */  
                background-color: #fff;  
                box-shadow: 0 1px 3px rgba(0,0,0,0.5); /* 添加阴影效果 */  
            }  
           QPushButton {  
                border: none;  
                border-radius: 5px;  
                padding: 10px 20px;  
                font-size: 12pt;  
                font-family: "Arial", sans-serif;  
                color: #fff;  
                background-color: #5cb85c; /* 更深的初始背景色 */  
                box-shadow: 0 4px 6px rgba(0, 0, 0, 0.15), 0 2px 4px rgba(0, 0, 0, 0.15) inset; /* 外部和内部阴影，增加立体感 */  
                transition: background-color 0.3s ease, box-shadow 0.3s ease;  
            }  
            
            QPushButton:hover {  
                background-color: #4cae4c; /* 鼠标悬停时背景色稍浅 */  
                box-shadow: 0 6px 10px rgba(0, 0, 0, 0.2), 0 3px 6px rgba(0, 0, 0, 0.2) inset; /* 增加阴影强度 */  
            }  
            
            QPushButton:pressed {  
                background-color: #449d44; /* 按下时背景色更深 */  
                box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2) inset, 0 1px 0 rgba(255, 255, 255, 0.1); /* 增加内部阴影和高光 */  
                padding-top: 11px; /* 稍微调整内边距以模拟按下效果 */  
                padding-bottom: 9px;  
            }
           """)
    def create_module_1(self):
        #第一个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        #创建第一个文本输入框并设置占位提示文字
        self.text_edit1 = QPlainTextEdit()
        self.text_edit1.setPlaceholderText("输入正向提示词")
        self.text_edit1.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        font_metrics = QFontMetrics(self.text_edit1.document().defaultFont())
        line_height = font_metrics.lineSpacing()

        visible_lines = 8  # 设置可见的行数
        desired_height = line_height * visible_lines
        self.text_edit1.setFixedHeight(desired_height)

        layout.addWidget(self.text_edit1)

        # 创建第二个文本输入框并设置占位提示文字  
        self.text_edit2 = QPlainTextEdit()
        self.text_edit2.setPlaceholderText("输入反向提示词")
        self.text_edit2.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        self.text_edit2.setFixedHeight(desired_height)
        layout.addWidget(self.text_edit2)
        # 创建水平布局用于放置滑块和标签  
        layout1 = QHBoxLayout()  
        layout1.setSpacing(5)  # 设置布局间距  
        layout1.setAlignment(Qt.AlignCenter)  # 设置布局对齐方式  
        
        # 创建滑块  
        self.width_pic = QSlider(Qt.Horizontal)  
        self.width_pic.setMinimum(0)  # 设置最小值  
        self.width_pic.setMaximum(512)  # 设置最大值  
        self.width_pic.setValue(256)  # 设置默认值（中心位置）  
        self.width_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置  
        self.width_pic.setTickInterval(10)  # 设置刻度间隔  
        self.width_pic.setStyleSheet("QSlider::groove:horizontal { border: 2px solid black; height: 10px; margin: 0; border-radius: 5px; } "  # 设置滑轨样式  
                                    "QSlider::handle:horizontal { background: black; border: 2px solid black; width: 15px; margin: -7px 0; border-radius: 8px; }")  # 设置滑块样式  
        
        # 创建标签用于显示滑块的值  
        self.label_width = QLabel('256')  
        self.label_width.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式  
        
        # 创建标签用于显示'width : '文本  
        lable_tmp = QLabel('width  : ')  
        lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示  
        
        # 将标签和滑块添加到水平布局中  
        layout1.addWidget(lable_tmp)  
        layout1.addWidget(self.width_pic)  
        layout1.addWidget(self.label_width)  
        
        # 将水平布局添加到垂直布局中（假设layout是已有的垂直布局）  
        layout.addLayout(layout1)  
        #创建标签用于显示'width : '文本  
       
##
        # 创建标签用于显示滑块的值
        #self.label_width = QLabel('256')
        #lable_tmp = QLabel('width  : ')

        # 将滑块的值改变信号与槽函数连接
        self.width_pic.valueChanged.connect(self.on_value_changed_width)
        layout1.addWidget(lable_tmp)
        layout1.addWidget(self.width_pic)
        layout1.addWidget(self.label_width)

        layout2 = QHBoxLayout()
##
      # 创建滑块  
        self.height_pic = QSlider(Qt.Horizontal)  
        self.height_pic.setMinimum(0)  # 设置最小值  
        self.height_pic.setMaximum(512)  # 设置最大值  
        self.height_pic.setValue(256)  # 设置默认值（中心位置）  
        self.height_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置  
        self.height_pic.setTickInterval(10)  # 设置刻度间隔  
        
        # 设置滑轨和滑块样式  
        self.height_pic.setStyleSheet("""  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """)  
        
        # 创建标签用于显示滑块的值  
        self.label_height = QLabel('256')  
        self.label_height.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式  
        
        # 创建标签用于显示'height : '文本  
        lable_tmp = QLabel('height : ')  
        lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示  
        
        # 假设layout1是一个已经存在的水平布局  
        layout2.addWidget(lable_tmp)  
        layout2.addWidget(self.height_pic)  
        layout2.addWidget(self.label_height)  
        
        # 假设layout是一个已经存在的垂直布局  
        # 将水平布局layout1添加到垂直布局layout中  
        layout.addLayout(layout2)  
        # 将滑块的值改变信号与槽函数连接  
        self.height_pic.valueChanged.connect(self.on_value_changed_height)  
        # 注意：下面这行代码不应该在这里，因为它会再次向layout1添加控件，这是重复的  
        #layout1.addWidget(self.label_height)    
        layout3 = QHBoxLayout()  
        
        # 如果需要更新标签的值，可以在on_value_changed_height槽函数中实现  
        def on_value_changed_height(self, value):  
            self.label_height.setText(str(value))  
        # 创建滑块  
        self.step_pic = QSlider(Qt.Horizontal)  
        self.step_pic.setMinimum(1)       # 设置最小值  
        self.step_pic.setMaximum(40)      # 设置最大值  
        self.step_pic.setValue(20)         # 设置默认值（中心位置）  
        self.step_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置  
        self.step_pic.setTickInterval(2)  # 设置刻度间隔  
        
        #设置滑轨和滑块样式  
        self.step_pic.setStyleSheet("""  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """)  
       
        # # 创建标签用于显示滑块的值  
        # self.label_step = QLabel('10')  
        # self.label_step.setAlignment(Qt.AlignCenter)  # 设置标签文本对齐方式 
        # # 创建标签用于显示'step : '文本  
        # lable_tmp = QLabel('step : ')  
        # lable_tmp.setAlignment(Qt.AlignRight)  # 设置标签文本对齐方式，使其靠右显示  
        # 创建标签用于显示滑块的值
        self.label_step = QLabel('20')
        lable_tmp = QLabel('step : ')
        # 创建一个垂直布局来包裹标签、滑块和显示值的标签  
        #layout3.addWidget(lable_tmp)  
        #layout3.addWidget(self.step_pic)  
        #layout3.addWidget(self.label_step)     
        # 将滑块的值改变信号与槽函数连接  
        self.step_pic.valueChanged.connect(self.on_value_changed_step)
        layout3.addWidget(lable_tmp)
        layout3.addWidget(self.step_pic)
        layout3.addWidget(self.label_step)

        layout.addLayout(layout1)
        layout.addLayout(layout2)
        layout.addLayout(layout3)
        widget.setLayout(layout)
        return widget
    
    def create_module_2(self):
        # 第二个模块的布局和内容
        widget = QWidget()
        layout = QVBoxLayout(widget)

        # 创建第一个文本输入框并设置占位提示文字
        self.text_edit3 = QPlainTextEdit()
        self.text_edit3.setPlaceholderText("输入嵌入数据")
        self.text_edit3.setStyleSheet("QPlainTextEdit { border: 2px solid black; }")

        font_metrics = QFontMetrics(self.text_edit3.document().defaultFont())
        line_height = font_metrics.lineSpacing()

        visible_lines = 8  # 设置可见的行数
        desired_height = line_height * visible_lines
        self.text_edit3.setFixedHeight(desired_height)

        layout.addWidget(self.text_edit3)  # 将文本输入框添加到布局中

        # 创建下拉框，并添加选项
        lay1 = QHBoxLayout()
        label = QLabel("差异范围：")  
        label.setFixedWidth(150)  # 设置固定宽度为120像素  
        lay1.addWidget(label)
        self.combo_box = QComboBox()
        self.combo_box.addItems(['0-7', '8-15', '16-31', '32-63', '64-127', '128-255'])
        lay1.addWidget(self.combo_box)  # 将下拉框添加到布局中
        layout.addLayout(lay1)

        # 设置下拉框的大小，这里可以根据实际情况调整  
        self.combo_box.setSizeAdjustPolicy(QComboBox.AdjustToContents)  
        font = self.combo_box.font()  
        font.setPointSize(15)  # 设置字体大小为15点  
        self.combo_box.setFont(font)  
        lay1.addWidget(self.combo_box)  # 将下拉框添加到布局中  
        layout.addLayout(lay1) 
        # 创建下拉框，并添加选项
        lay2 = QHBoxLayout()
        label = QLabel("checkpoint：")
        lay2.addWidget(label)
        self.combo_box1 = QComboBox()
        self.combo_box1.addItems(["3Guofeng3_v34.safetensors", "blazingrealdrive_V02h.safetensors", "v2-1_768-ema-pruned.ckpt"])
        # 设置下拉框选项的字体大小  
        font = self.combo_box1.font() 
        font.setPointSize(12)  # 设置字体大小为12点，你可以根据需要调整这个值  
        self.combo_box1.setFont(font) 
        lay2.addWidget(self.combo_box1)  # 将下拉框添加到布局中
        layout.addLayout(lay2)

        layout3 = QHBoxLayout()

        layout2 = QHBoxLayout()
        label = QLabel("请输入保存文件名：")
        layout2.addWidget(label)
        self.pic_name = QLineEdit()
        self.pic_name.setStyleSheet("QLineEdit { border: 2px solid black; }")
        layout2.addWidget(self.pic_name)
        layout.addLayout(layout2)

        # 创建滑块
        self.CFG_pic = QSlider(Qt.Horizontal)
        self.CFG_pic.setMinimum(1)  # 设置最小值
        self.CFG_pic.setMaximum(10)  # 设置最大值
        self.CFG_pic.setValue(5)  # 设置默认值（中心位置）
        self.CFG_pic.setTickPosition(QSlider.TicksBelow)  # 设置刻度位置
        self.CFG_pic.setTickInterval(1)  # 设置刻度间隔
        #设置滑轨和滑块样式  
        self.CFG_pic.setStyleSheet("""  
            QSlider::groove:horizontal {  
                border: 2px solid black;  
                height: 10px;  
                margin: 0;  
                border-radius: 5px;  
            }  
            QSlider::handle:horizontal {  
                background: black;  
                border: 2px solid black;  
                width: 15px;  
                margin: -7px 0;  
                border-radius: 8px;  
            }  
        """)  
        # 创建标签用于显示滑块的值
        self.label_CFG = QLabel('5')
        lable_tmp = QLabel('CFG Scale : ')
        # 将滑块的值改变信号与槽函数连接
        self.CFG_pic.valueChanged.connect(self.on_value_changed_CFG)
        layout3.addWidget(lable_tmp)
        layout3.addWidget(self.CFG_pic)
        layout3.addWidget(self.label_CFG)
        layout.addLayout(layout3)
        
        # 创建两个按钮，并排对称排布
        button1 = QPushButton("生成图片")
        button2 = QPushButton("嵌入水印")
        button3 = QPushButton("水印提取")
        button1.clicked.connect(self.on_button_clicked1)
        button2.clicked.connect(self.on_button_clicked2)
        button3.clicked.connect(self.on_button_clicked3)
        # central_widget = QtWidgets.QWidget()
        # self.page2 = Page2(central_widget)


        # 为了使两个按钮对称排布，我们可以将它们放入一个水平布局中，然后将这个布局添加到主布局中
        sub_layout = QHBoxLayout()
        sub_layout.addWidget(button1)
        sub_layout.addWidget(button2)
        sub_layout.addWidget(button3)
        sub_layout.setAlignment(Qt.AlignCenter)  # 设置按钮居中对齐

        layout.addLayout(sub_layout)  # 将按钮布局添加到主布局中

        widget.setLayout(layout)
        return widget

    def create_module_3(self):
        # 第三个模块的布局和内容
        image_path = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\SD\\LN.jpg'
        widget = QWidget()
        layout = QHBoxLayout(widget)

        # 创建一个QLabel对象用于显示图片
        self.image_label_init = QLabel()

        # 使用QPixmap加载图片
        pixmap = QtGui.QPixmap(image_path)

        # 如果图片大小大于标签当前大小，则调整标签大小以适应图片
        self.image_label_init.setPixmap(pixmap.scaled(self.image_label_init.size(), aspectRatioMode=Qt.KeepAspectRatio))

        # 将QLabel添加到布局中
        layout.addWidget(self.image_label_init)

        # 设置布局到widget
        widget.setLayout(layout)
        return widget

    def create_module_4(self):
        # 第三个模块的布局和内容
        image_path = 'C:\\Users\\chengwenjun\\Desktop\\stablediffusion\\SD\\LN1.png'
        widget = QWidget()
        layout = QHBoxLayout(widget)

        # 创建一个QLabel对象用于显示图片
        self.image_label_ans = QLabel()

        # 使用QPixmap加载图片
        pixmap = QtGui.QPixmap(image_path)

        # 如果图片大小大于标签当前大小，则调整标签大小以适应图片
        self.image_label_ans.setPixmap(pixmap.scaled(self.image_label_ans.size(), aspectRatioMode=Qt.KeepAspectRatio))

        # 将QLabel添加到布局中
        layout.addWidget(self.image_label_ans)

        # 设置布局到widget
        widget.setLayout(layout)
        return widget

    def on_value_changed_width(self, value):
        # 更新标签显示的值
        self.label_width.setText(str(value))

    def on_value_changed_height(self, value):
        # 更新标签显示的值
        self.label_height.setText(str(value))

    def on_value_changed_step(self, value):
        # 更新标签显示的值
        self.label_step.setText(str(value))

    def on_value_changed_CFG(self, value):
        # 更新标签显示的值
        self.label_CFG.setText(str(value))

    def on_button_clicked3(self):
        page22 = Page2(self)
        page22.show()


    def on_button_clicked1(self):
        width = self.width_pic.value()
        heigth = self.height_pic.value()
        step = self.step_pic.value()
        cfg = self.CFG_pic.value()
        negative_prompt = self.text_edit2.toPlainText()
        prompt = self.text_edit1.toPlainText()
        checkpoint = self.combo_box1.currentText()
        if not prompt and not negative_prompt:
            # 如果为空，显示提示框
            QMessageBox.warning(self, '警告', '正向/反向提示词不能同时为空')
            # 退出函数，不再执行后续代码
            return
        payload = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "steps": step,
            "width": width,
            "height": heigth,
            "seed": -1,
            "CFG Scale": cfg,
            "stable diffusion checkpoint": checkpoint,  # 下拉框
        }
        self.get_pic = 1
        # 更换图片
        # 使用Pillow打开图片
        self.image_path = self.pic_name.text()
        if len(self.image_path) == 0:
            print('路径为空！不符合要求')
            #self.image_path = 'example.png'
        else:
            self.image_path = self.image_path + '.png'
            #print(payload)
        print(self.image_path)
        get_pic_from_stable_diffusion(payload,self.image_path)
        pil_image = Image.open(self.image_path)
        try:
            # 将Pillow图像转换为QImage
            qimage = QImage(pil_image.tobytes(), pil_image.size[0], pil_image.size[1], QImage.Format.Format_RGB888)
        except Exception as e:
            print("An error occurred:", e)
        # 将QImage转换为QPixmap
        try:
            pixmap = QPixmap.fromImage(qimage)
        except Exception as e:
            print("An error occurred:", e)
        # 在QLabel中显示灰度图片
        self.image_label_init.setPixmap(pixmap)
        # 在QLabel中显示灰度图片
        self.image_label_init.setPixmap(pixmap)

    def on_button_clicked2(self):
        data = self.text_edit3.toPlainText()
        choose = self.combo_box.currentText()
        choose = choose.split('-')
        min_d = int(choose[0])
        max_d = int(choose[1])
        fw = (min_d,max_d)
        if len(data) == 0:
            QMessageBox.warning(self, '警告', '嵌入数据不可为空')
            return
        if not os.path.exists(self.image_path):
            QMessageBox.warning(self, '警告', '图片尚未生成')
            return
        l = actor(self.image_path,data,fw)
        QMessageBox.warning(self,'提示','加密长度为： '+str(l))
        # 使用Pillow打开图片
        self.image_path = 'ans.png'
        pil_image = Image.open(self.image_path)

        try:
            # 将Pillow图像转换为QImage
            qimage = QImage(pil_image.tobytes(), pil_image.size[0], pil_image.size[1], QImage.Format.Format_RGB888)
        except Exception as e:
            print("An error occurred:", e)
        # 将QImage转换为QPixmap
        try:
            pixmap = QPixmap.fromImage(qimage)
        except Exception as e:
            print("An error occurred:", e)
        # 在QLabel中显示灰度图片
        self.image_label_ans.setPixmap(pixmap)
        # 在QLabel中显示灰度图片
        self.image_label_ans.setPixmap(pixmap)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    main_window = Page1()
    main_window.show()
    sys.exit(app.exec_())